#coding=utf-8

import numpy, scipy
import matplotlib.pyplot as plt
from numpy import ma
from scipy import sparse
import netCDF4

import os, sys
import shutil, gzip

from mpl_toolkits.basemap import Basemap, cm

class data_map:

    def __init__(self, ds_path, workspace=None, hurricane_name=None, height=None, date_time=None):
        print "WORKSPACE:", os.environ["TEMP"]
        self.workspace = workspace if workspace != None else os.environ["TEMP"]
        self.archive = (hurricane_name + "_" + height + "_" + date_time).replace("-","_").replace(" ", "")
        ds_out = os.path.join(self.workspace, os.path.basename(ds_path)).rstrip('.gz')
        #Unzip dataset to workspace
        f_in = gzip.open(ds_path, 'rb')
        f_out = open(ds_out, 'wb')
        f_out.write(f_in.read())
        f_out.close()
        f_in.close()
        self.Dataset = netCDF4.Dataset(ds_out)
        pass

    def __draw_core(self, base_map, m_lat, m_lon, z, zm, no_save=False):
        g_lon, g_lat = numpy.meshgrid(m_lon, m_lat)
        prj_lon, prj_lat = base_map(g_lon, g_lat)
        #Save archive
        if not no_save:
            archive_path = os.path.join(self.workspace, self.archive)
            numpy.savez_compressed(archive_path, X=g_lon, Y=g_lat, Z=z)
        levels = numpy.arange(0, numpy.max(z), 2.0).tolist()
        base_map.contourf(prj_lon, prj_lat, zm, levels=levels, cmap=cm.s3pcpn)

    def draw_basic(self, base_map):
        m_lat = self.Dataset.variables['Lat'][:]
        m_lon = self.Dataset.variables['Lon'][:]
        z = numpy.array(self.Dataset.variables['MergedReflectivityQC'][0])
        assert len(z.shape) == 2
        #Mask value less than 1
        zm = ma.masked_less(z, 1)
        self.__draw_core(base_map, m_lat, m_lon, z, zm)
        pass

    def draw_sparse(self, base_map):
        """Draw a map from SparseLatLonGrid"""
        #Get dimensions and range
        row_count = len(self.Dataset.dimensions['Lat'])
        col_count = len(self.Dataset.dimensions['Lon'])
        #Convert sparse matrix to dense array
        pixel_x = self.Dataset.variables['pixel_x']
        pixel_y = self.Dataset.variables['pixel_y']
        reflect = self.Dataset.variables['MergedReflectivityQC']
        #Encoding a sparse matrix is slow, directly fill it!
        #reflect_array = numpy.array(sparse.csr_matrix((reflect, (pixel_x, pixel_y)), shape=(row_count, col_count), dtype=numpy.float32).todense())
        reflect_array = numpy.empty((row_count, col_count))
        for i in xrange(len(pixel_x)):
            reflect_array[pixel_x[i], pixel_y[i]] = reflect
        #Mask value less than 1
        masked_reflect = ma.masked_less(reflect_array, 1)
        #Produce a LonLatGrid meshgrid
        m_lon = numpy.arange(self.Dataset.Longitude, self.Dataset.Longitude + col_count * self.Dataset.LonGridSpacing, self.Dataset.LonGridSpacing)
        m_lat = numpy.arange(self.Dataset.Latitude, self.Dataset.Latitude - row_count * self.Dataset.LatGridSpacing, -self.Dataset.LatGridSpacing)
        self.__draw_core(base_map, m_lat, m_lon, reflect_array, masked_reflect)
        pass


    def draw_archive(self, base_map, npz_path):
        data = numpy.load(npz_path)
        g_lon = data['X']
        g_lat = data['Y']
        prj_lon, prj_lat = base_map(g_lon, g_lat)
        value = data['Z']
        zm = ma.masked_less(value, 0)
        levels = numpy.arange(0, numpy.max(value), 2.0).tolist()
        base_map.contourf(prj_lon, prj_lat, zm, levels, cmap=cm.s3pcpn)
        pass

    def draw(self):
        """Draw map"""
        m = Basemap(width=12000000,height=9000000, rsphere=(6378137.00,6356752.3142), resolution='l', area_thresh=1000.0, projection='lcc', lat_1=45.0,lat_2=55,lat_0=50,lon_0=-107.0)
        #m.fillcontinents(color='coral',lake_color='aqua')
        m.drawcountries()
        m.drawstates()
        m.drawcoastlines()
        # draw parallels and meridians.
        m.drawparallels(numpy.arange(-80.,81.,20.))
        m.drawmeridians(numpy.arange(-180.,181.,20.))
        m.drawmapboundary(fill_color='white')
        archive_path = os.path.join(self.workspace, self.archive)
        if os.path.exists(archive_path):
            self.draw_archive(m, archive_path)
        elif os.path.exists(archive_path + ".npz"):
            self.draw_archive(m, archive_path + ".npz")
        elif self.Dataset.DataType == 'LatLonGrid':
            self.draw_basic(m)
        else:
            self.draw_sparse(m)
        plt.colorbar()
        plt.show()
        plt.clf()
        del m

